use crate::error::{Error, Result};
use crate::protocol::varint::{read_varint, write_varint};
use crate::protocol::TCP_REQUEST_ID;
use super::padding::{TCP_REQUEST_PADDING, TCP_RESPONSE_PADDING};

#[derive(Debug, Clone)]
pub struct TcpRequest {
    pub address: String,
    pub padding: Vec<u8>,
}

impl TcpRequest {
    pub fn new(address: String) -> Self {
        Self {
            address,
            padding: Vec::new(),
        }
    }

    pub fn encode(&self) -> Vec<u8> {
        let mut buf = Vec::new();

        // Frame type
        buf.extend_from_slice(&encode_varint(0x401));

        // Write request ID (varint)
        write_varint(&mut buf, TCP_REQUEST_ID);

        // Write address length and address
        write_varint(&mut buf, self.address.len() as u64);
        buf.extend_from_slice(self.address.as_bytes());

        // Padding
        let padding = TCP_REQUEST_PADDING.generate();
        buf.extend_from_slice(&encode_varint(padding.len() as u64));
        buf.extend_from_slice(&padding);

        buf
    }

    pub fn decode(data: &[u8]) -> Result<Self> {
        let mut cursor = 0;

        // Read frame type
        let (frame_type, n) = read_varint(&data[cursor..])?;
        cursor += n;
        if frame_type != 0x401 {
            return Err(Error::InvalidFrameType);
        }

        // Read request ID
        let (id, n) = read_varint(&data[cursor..])?;
        cursor += n;
        if id != TCP_REQUEST_ID {
            return Err(Error::InvalidRequestId);
        }

        // Read address
        let (addr_len, n) = read_varint(&data[cursor..])?;
        cursor += n;
        let address = String::from_utf8(data[cursor..cursor + addr_len as usize].to_vec())?;
        cursor += addr_len as usize;

        // Read padding
        let (pad_len, n) = read_varint(&data[cursor..])?;
        cursor += n;
        let padding = data[cursor..cursor + pad_len as usize].to_vec();

        Ok(Self { address, padding })
    }
}

#[derive(Debug, Clone)]
pub struct TcpResponse {
    pub status: u8,
    pub message: String,
    pub padding: Vec<u8>,
}

impl TcpResponse {
    pub fn ok() -> Self {
        Self {
            status: 0x00,
            message: "OK".to_string(),
            padding: Vec::new(),
        }
    }

    pub fn error(message: String) -> Self {
        Self {
            status: 0x01,
            message,
            padding: Vec::new(),
        }
    }

    pub fn is_ok(&self) -> bool {
        self.status == 0x00
    }

    pub fn encode(&self) -> Vec<u8> {
        let mut buf = Vec::new();

        // Status
        buf.push(if self.is_ok() { 0 } else { 1 });

        // Message length and message
        let msg_bytes = self.message.as_bytes();
        buf.extend_from_slice(&encode_varint(msg_bytes.len() as u64));
        buf.extend_from_slice(msg_bytes);

        // Padding
        let padding = TCP_RESPONSE_PADDING.generate();
        buf.extend_from_slice(&encode_varint(padding.len() as u64));
        buf.extend_from_slice(&padding);

        buf
    }

    pub fn decode(data: &[u8]) -> Result<Self> {
        let mut cursor = 0;

        // Read status
        let status = data[cursor];
        cursor += 1;

        let (msg_len, n) = read_varint(&data[cursor..])?;
        cursor += n;
        let message = String::from_utf8(data[cursor..cursor + msg_len as usize].to_vec())?;
        cursor += msg_len as usize;

        let (pad_len, n) = read_varint(&data[cursor..])?;
        cursor += n;
        let padding = data[cursor..cursor + pad_len as usize].to_vec();

        Ok(Self {
            status,
            message,
            padding,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_tcp_request_roundtrip() {
        let req = TcpRequest::new("example.com:443".to_string());
        let encoded = req.encode();
        let decoded = TcpRequest::decode(&encoded).unwrap();
        assert_eq!(decoded.address, "example.com:443");
    }

    #[test]
    fn test_tcp_response_roundtrip() {
        let resp = TcpResponse::ok();
        let encoded = resp.encode();
        let decoded = TcpResponse::decode(&encoded).unwrap();
        assert!(decoded.is_ok());
        assert_eq!(decoded.message, "OK");
    }
}
